import os
import random
import glob
from PIL import Image, ImageEnhance, ImageOps

DESKTOP_PATH = os.path.join(os.path.expanduser("~"), "Desktop")
OUTPUT_DIR = os.path.join(DESKTOP_PATH, "CNN_model")
ORIGINAL_DIR = os.path.join(OUTPUT_DIR, "phone_screens_original")
NORMAL_IMAGES_DIR = os.path.join(ORIGINAL_DIR, "normal_screens")  # 原始正常图片目录
ENHANCED_NORMAL_DIR = os.path.join(OUTPUT_DIR,"CNN_test","normal")  # 增强后的正常图片目录

IMAGES_PER_ORIGINAL = 10  # 每张原始图片生成的增强图片数量

# 确保输出目录存在
os.makedirs(ENHANCED_NORMAL_DIR, exist_ok=True)

def load_images_from_dir(directory, extensions=('png','jpg')):
    """从目录加载所有指定格式的图片"""
    images = []
    for ext in extensions:
        for img_path in glob.glob(os.path.join(directory, f'*.{ext}')):
            try:
                img = Image.open(img_path).convert("RGBA")
                images.append((img, os.path.basename(img_path)))
                print(f"加载图片: {img_path}")
            except Exception as e:
                print(f"加载图片失败 {img_path}: {e}")
    return images

def random_enhance(image):
    """对图片进行随机微调增强"""
    img = image.copy()
    
    # 随机放缩 (±5%)
    scale_factor = random.uniform(0.95, 1.05)
    new_width = int(img.width * scale_factor)
    new_height = int(img.height * scale_factor)
    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
    
    # 随机旋转 (-3° 到 3°)
    rotate_angle = random.uniform(-3, 3)
    img = img.rotate(rotate_angle, expand=True, fillcolor=(0, 0, 0, 0))
    
    # 随机裁剪，确保尺寸不小于原图的90%
    crop_ratio = random.uniform(0.9, 1.0)
    new_width = int(img.width * crop_ratio)
    new_height = int(img.height * crop_ratio)
    
    if new_width < img.width and new_height < img.height:
        left = random.randint(0, img.width - new_width)
        top = random.randint(0, img.height - new_height)
        img = img.crop((left, top, left + new_width, top + new_height))
    
    # 随机亮度调整 (±10%)
    brightness_factor = random.uniform(0.9, 1.1)
    enhancer = ImageEnhance.Brightness(img)
    img = enhancer.enhance(brightness_factor)
    
    # 随机对比度调整 (±10%)
    contrast_factor = random.uniform(0.9, 1.1)
    enhancer = ImageEnhance.Contrast(img)
    img = enhancer.enhance(contrast_factor)
    
    return img

def generate_enhanced_normal_images(original_images):
    """生成增强的正常图片数据集"""
    enhanced_images = []
    for img, img_name in original_images:
        base_name = os.path.splitext(img_name)[0]
        
        for i in range(IMAGES_PER_ORIGINAL):
            enhanced_img = random_enhance(img)
            output_path = os.path.join(ENHANCED_NORMAL_DIR, f"{base_name}_enhanced_{i}.png")
            enhanced_img.convert("RGB").save(output_path)  # 转换为RGB保存
            enhanced_images.append((enhanced_img, f"{base_name}_enhanced_{i}.png"))
            print(f"生成增强图片: {output_path}")
    
    return enhanced_images

def main():
    # 加载原始图片
    print("加载原始正常图片...")
    original_normal_images = load_images_from_dir(NORMAL_IMAGES_DIR)
    
    if not original_normal_images:
        print("错误: 没有找到原始正常图片，请检查NORMAL_IMAGES_DIR目录")
        return
    
    
    # 生成增强的正常图片
    print("开始生成增强的正常图片...")
    enhanced_normal_images = generate_enhanced_normal_images(original_normal_images)
    print(f"已生成 {len(enhanced_normal_images)} 张增强的正常图片")
    
    print("数据集生成完成!")

if __name__ == "__main__":
    main()
